import os
import re
import sys
import cv2 as cv
import numpy as np
from PyCmpltrtok.common import rand_color
from PyCmpltrtok.common_opencv import imzoom_rect

NUM = 2
TGT_RECT = (1500 // NUM, 700)

voc2012test_list_path = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/data/voc.2012_mini.test'
result_dir = '/var/asuspei/my_svn/c_darknet/darknet-debug-cpu/results/voc.2012_mini.test'

# get all img paths
img_paths = []
with open(voc2012test_list_path, 'r', encoding='utf8') as f:
    while True:
        xpath = f.readline()
        if not xpath:
            break
        img_paths.append(xpath)
# print(img_paths)
for i, img_path in enumerate(img_paths):
    if img_path[-2:] == '\r\n':
        img_paths[i] = img_path[:-2]
        continue
    elif img_path[-1:] == '\n':
        img_paths[i] = img_path[:-1]
        continue
# print(img_paths)

# map main name to path and to boxes
name_list = []
name2path, name2boxes = {}, {}
for img_path in img_paths:
    dir_name, base_name = os.path.split(img_path)
    main_name, ext_name = os.path.splitext(base_name)
    name_list.append(main_name)
    name2path[main_name] = img_path
    name2boxes[main_name] = []

# read in box info
result_file_names = os.listdir(result_dir)
regexp = re.compile(r'_([^\._]+)\.txt$')
for result_file_name in result_file_names:
    result_file_path = os.path.join(result_dir, result_file_name)
    matcher = regexp.search(result_file_name)
    if not matcher:
        print(f'File {result_file_name} named badly!', file=sys.stderr)
        continue
    label = matcher.group(1)
    boxes = []
    with open(result_file_path, 'r') as f:
        while True:
            line = f.readline()
            if not line:
                break
            vals = line.split()
            main_name = vals[0]
            objness = vals[1]
            box_info = [label, float(objness)]
            box_info.extend([int(np.round(float(x))) for x in vals[2:]])
            name2boxes[main_name].append(box_info)

xpos = 0
xlen = len(name_list)
while True:
    xnames = name_list[xpos:xpos+NUM]
    img_arr = []
    for xname in xnames:
        xpath = name2path[xname]
        xboxes = name2boxes[xname]
        img = cv.imread(xpath, cv.IMREAD_COLOR)
        for box in xboxes:
            label, objness, x1, y1, x2, y2 = box
            color = rand_color()
            cv.rectangle(img, (x1, y1), (x2, y2), color, 2)
            cv.putText(img, f'{label} ({objness:.4f})', (x1, y1 + 15), cv.FONT_HERSHEY_PLAIN, 1, color, 1)
        img = imzoom_rect(img, TGT_RECT)
        img_arr.append(img)
    img = np.concatenate(img_arr, axis=1)
    cv.line(img, (TGT_RECT[0], 1), TGT_RECT, (0, 255, 0), 2)
    cv.imshow('check YOLO valid', img)
    k = cv.waitKey()
    if ord('p') == k:
        xpos -= 1
    elif ord('n') == k:
        xpos += 1
    elif ord('q') == k:
        print('Bye!')
        break
    xpos = xpos if xpos >= 0 else 0
    xpos = xpos if xpos < xlen else xlen - 1
    cv.destroyAllWindows()
